# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime
import logging

import numpy

from tornado.ioloop import PeriodicCallback

from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp
from coin.strategy.mm.base.feed import need_check_limit_price


class FeedArbSubscriberHandler(object):
  SubscriberType = None

  def __init__(self, *, worker_id, ioloop, queue, native_symbols, logger=None):
    self._queue = queue
    self._native_symbols = native_symbols
    self._worker_id = worker_id
    self._ioloop = ioloop
    self._logger = logger or logging.getLogger(__name__)
    self._subscriber = None
    self._available = False
    self._tsdiff = []
    # sorry. super stupid hack..
    # beware bitmex.. better not restart for bitmex...
    self._allow_reconnect = str(type(self)).find("Bitmex") < 0

  def init_subscriber(self):
    if self._subscriber is not None:
      self._subscriber.stop()
      self._logger.info('Restarting %s subscriber...' % self._worker_id)

    self._subscriber = self.SubscriberType(self._queue,
                                           'a',
                                           self._worker_id,
                                           self._native_symbols,
                                           close_callback=self._on_disconnect)
    self._subscriber.start(self._ioloop)
    self._available = True

  @property
  def available(self):
    return self._available

  def _on_disconnect(self):
    self._available = False
    self._logger.error('%s Disconnected. Restarting in 5 seconds.' % self._worker_id)
    self._ioloop.add_timeout(datetime.timedelta(seconds=5), self.init_subscriber)

  def restart(self):
    try:
      if self._subscriber is not None:
        if self._allow_reconnect:
          self._subscriber.stop(trigger_close_callback=True)
        else:
          self._subscriber.stop(trigger_close_callback=False)
        return True
    except Exception:
      self._logger.exception('Error stopping %s subscriber' % self._worker_id)
      self._on_disconnect()
      return False

  def stop(self):
    try:
      if self._subscriber is not None:
        self._subscriber.stop()
        return True
    except Exception:
      self._logger.exception('Error stopping %s subscriber' % self._worker_id)
      return False

  def monitor_feed(self, arb_book_builder):
    last_feed_ts = arb_book_builder.get_worker_last_feed_ts(self._worker_id)
    if last_feed_ts is None:
      return

    diff = get_timestamp() - last_feed_ts
    self._logger.info('Last feed: %d ms ago' % (diff / 10**6))

    if hasattr(self._subscriber, "BookFeedDisconnectWindowInNs"):
      conn_window = self._subscriber.BookFeedDisconnectWindowInNs
    else:
      conn_window = 2 * 60 * 1e9  # 2 mins

    if diff < conn_window:
      return
    else:
      self._logger.error('Monitor feed detect discontinuity. thres: %d(ms)', (conn_window * 1e-6))

    try:
      if self._subscriber is not None:
        self._subscriber.stop()
    except Exception:
      self._logger.exception('Error stopping %s subscriber' % self._worker_id)
    self._on_disconnect()

  def record_tsdiff(self, timestamp, arb_book_builder):
    last_feed_ts = arb_book_builder.get_worker_last_feed_ts(self._worker_id)
    if last_feed_ts is None:
      return
    self._tsdiff.append(timestamp - last_feed_ts)
    self._tsdiff = self._tsdiff[-10:]

  def get_mean_tsdiff(self):
    if len(self._tsdiff) == 0:
      return None
    return numpy.median(self._tsdiff)


class FeedArbSubsystem(object):
  FlowBookBuilder = None
  FeedSubscriber = None
  FeedSubReq = None
  _cooltime_reconnect = datetime.timedelta(seconds=3)

  def __init__(self,
               product,
               on_book_init_callback,
               worker_ids=None,
               feed_checker_config=None,
               fs_config=None,
               logger=None):
    self._logger = logger or logging.getLogger(__name__)

    self._product = to_list(product)
    self._on_book_init_callback = on_book_init_callback

    feedsub_str = self.FeedSubReq.to_str()
    if worker_ids is not None:
      self._worker_ids = list(worker_ids)
    else:
      if fs_config is None:
        num_feed_worker = 2
      elif fs_config.feedsub_arbitration_override.get(feedsub_str, None) is not None:
        num_feed_worker = fs_config.feedsub_arbitration_override.get(feedsub_str, None)
      else:
        num_feed_worker = fs_config.arbitration_connection_count
      self._worker_ids = [(wid + 1) for wid in range(num_feed_worker)]

    assert len(self._worker_ids) > 1
    self._ioloop = None
    self._queue = None
    self._flow_sub = None

    self._subscriber_handler = {}
    self._book_builder = None

    self._available = True
    self._fs_config = fs_config

    self._feed_checker_config = feed_checker_config
    self._feed_monitor_periodic_callback = None
    self._check_limit_price = need_check_limit_price(self.FeedSubReq, fs_config)
    self._price_limit_error_callback = None

  @property
  def available(self):
    if not self._available:
      return False
    for worker_id in self._worker_ids:
      if self._subscriber_handler.get(worker_id, None):
        if self._subscriber_handler[worker_id].available:
          return True
    return False

  def set_price_limit_error_callback(self, callback):
    raise NotImplementedError(self.FeedSubReq)

  def _record_latency_agent(self):
    if self._book_builder is None:
      return
    timestamp = get_timestamp()
    for worker_id in self._worker_ids:
      worker = self._subscriber_handler.get(worker_id, None)
      if worker and worker.available:
        worker.record_tsdiff(timestamp, self._book_builder)

  def _switch_slow_agent(self):
    if self._book_builder is None:
      return
    tsdiff_max = None
    slow_worker_id = None
    num_avail_worker = 0
    for worker_id in self._worker_ids:
      worker = self._subscriber_handler.get(worker_id, None)
      if worker and worker.available:
        tsdiff = worker.get_mean_tsdiff()
        self._logger.info("worker %s: avg tds: %s", worker_id, tsdiff)
        if tsdiff is None:
          continue
        num_avail_worker += 1
        if (tsdiff_max is None) or (tsdiff > tsdiff_max):
          tsdiff_max = tsdiff
          slow_worker_id = worker_id
    if slow_worker_id is not None and num_avail_worker > 1:
      self._logger.info("restart worker %s: avg tds: %s" % (slow_worker_id, tsdiff_max))
      self._subscriber_handler.get(slow_worker_id, None).restart()

  def _monitor_feed(self):
    if self._book_builder is None:
      return
    for worker_id in self._worker_ids:
      if self._subscriber_handler.get(worker_id, None):
        self._subscriber_handler[worker_id].monitor_feed(self._book_builder)

  def _monitor_feed_checker(self):
    assert self._feed_checker is not None

    timestamp = get_timestamp()
    state = self._feed_checker.check_state(timestamp=timestamp)
    if state.ready:
      self._available = True
      return
    if not state.invalid:
      return

    self._feed_checker.print_status(timestamp)
    self._feed_checker.restart(timestamp=timestamp)
    self._logger.error('Feed checker result is invalid, reconnect...')

    try:
      self.stop()
    except Exception:
      self._logger.exception('Error stopping subscriber')
    self._on_disconnect()

  def _monitor_limit_price_checker(self):
    state = self._limit_price_checker.check_state()
    if not state:
      self._limit_price_checker.print_status()

    if self._price_limit_error_callback is not None:
      self._price_limit_error_callback(state)

  def _create_book_builder(self):
    if self._check_limit_price is not None:
      self._book_builder = self.FlowBookBuilder(self._flow_sub,
                                                self._worker_ids,
                                                self._feed_checker_config,
                                                check_limit_price=self._check_limit_price)
    else:
      self._book_builder = self.FlowBookBuilder(self._flow_sub,
                                                self._worker_ids,
                                                self._feed_checker_config)
    self._book_builder.set_live(True)

  def _init_subscriber(self):
    # Initialize subscriber
    native_symbols = [product.native_subscription_symbol for product in self._product]
    for worker_id in self._worker_ids:
      self._subscriber_handler[worker_id] = self.FeedSubscriber(worker_id=worker_id,
                                                                ioloop=self._ioloop,
                                                                queue=self._queue,
                                                                native_symbols=native_symbols)
      self._subscriber_handler[worker_id].init_subscriber()

    # Initialize book builder
    self._create_book_builder()

    # Callback
    self._on_book_init_callback(self.FeedSubReq, self._book_builder)

    # Feed checker
    if self._feed_checker_config is not None:
      self._feed_checker = self._book_builder._feed_handler._handler.feed_checker
      if self._feed_checker:
        self._feed_checker.restart(timestamp=get_timestamp())

    self._available = True

  def _on_disconnect(self):
    self._available = False
    self._logger.error('Disconnected. Restarting in %s.', self._cooltime_reconnect)
    self._ioloop.add_timeout(self._cooltime_reconnect, self._init_subscriber)

  def start_agent_switch(self, ioloop, switch_slow_agent_frequency_sec):
    assert switch_slow_agent_frequency_sec > 15, switch_slow_agent_frequency_sec
    self._record_latency_callback = PeriodicCallback(self._record_latency_agent,
                                                     switch_slow_agent_frequency_sec * 0.05 * 1000)
    self._record_latency_callback.start()
    self._agent_switch_callback = PeriodicCallback(self._switch_slow_agent,
                                                   switch_slow_agent_frequency_sec * 1000)
    self._agent_switch_callback.start()

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._queue = queue
    self._flow_sub = flow_sub

    # Init feed subscriber
    self._init_subscriber()

    # Setup feed monitor
    self._feed_monitor_periodic_callback = PeriodicCallback(self._monitor_feed, 60 * 1000)
    self._feed_monitor_periodic_callback.start()

    # Setup feed checker monitor
    if self._feed_checker_config is not None:
      self._feed_checker_monitor_periodic_callback = PeriodicCallback(self._monitor_feed_checker,
                                                                      1 * 1000)
      self._feed_checker_monitor_periodic_callback.start()

    # Setup limit price checker monitor
    if self._check_limit_price:
      self._limit_price_checker = \
          self._book_builder._feed_handler._handler.limit_price_checker
      self._limit_price_checker_periodic_callback = PeriodicCallback(
          self._monitor_limit_price_checker, 5 * 1000)
      self._limit_price_checker_periodic_callback.start()

  def stop(self):
    for worker_id in self._worker_ids:
      if self._subscriber_handler.get(worker_id, None) is not None:
        self._subscriber_handler[worker_id].stop()

  def restart(self):
    for worker_id in self._worker_ids:
      if self._subscriber_handler.get(worker_id, None) is not None:
        self._subscriber_handler[worker_id].restart()
